import datetime
import math
import random

import numpy as np
from treelib import Tree
from config.parameter_carrier import ParameterCarrier
from data_preparation.trajectory_set import TrajectorySet
from discretization.grid import Grid
from tools.noise import Noise


class TreeModel:
    def __init__(self, cc: ParameterCarrier):
        self.cc = cc
        # self.max_depth = cc.max_tree_depth
        self.max_depth = 5
        self.neighbors = []
        self.subcell_number = -1
        self.threshold = 10
        self.grid = Grid(self.cc)
        self.level1_cell_number = -1
        self.tree = None
        self.false_nodes = []
        self.children_all_zero_nodes = []

    def init_tree_constructor(self, grid: Grid, trajectory_set1: TrajectorySet):
        cc1 = self.cc
        # print("轨迹长度阈值：", trajectory_set1.lmax)
        # self.threshold = (trajectory_set1.lmax * math.log(trajectory_set1.trajectory_number / 2)) / (cc1.total_epsilon * (cc1.epsilon_partition[1] + cc1.epsilon_partition[2]))
        # print("阈值 ：",self.threshold)
        self.neighbors = grid.subcell_neighbors_real_level2_index
        self.subcell_number = grid.subcell_number
        self.level1_cell_number = grid.level1_cell_number
        tree = Tree()
        tree.create_node("root", "root", data={"value": 0, "path": [-1], "leaf": False})
        self.bulid_tree(tree, grid, trajectory_set1)


    ###tag存放节点对应的网格编号
    ###identifier存放节点的唯一标识符，不能重复，所以存放的是层数+这个节点在这一层的序号
    def bulid_tree(self, tree, grid: Grid, trajectory_set1: TrajectorySet):
        cc1 = self.cc
        total_epsilon = cc1.total_epsilon * (cc1.epsilon_partition[1] + cc1.epsilon_partition[2])
        max_tree_depth = self.max_depth
        print("树的最大深度：",max_tree_depth)
        print("树的阈值：",self.threshold)
        # print("树的第一层节点数：",self.subcell_number)#141
        current_depth = 1
        sensitivity = 1
        last_depth_only_leaf_nodes = False
        level_fisrt_data = grid.level_fisrt_node_value  # numpy
        noise1 = Noise()
        grid.level_fisrt_node_value_noisy = noise1.add_laplace_for_tree(level_fisrt_data, total_epsilon  / max_tree_depth, sensitivity)
        remain_epsilon = total_epsilon * (1 - 1 / max_tree_depth)
        print("开始构造树")
        print(datetime.datetime.now())
        while not last_depth_only_leaf_nodes and current_depth <= max_tree_depth:
            if current_depth == 1:
                for layer_index in range(self.subcell_number):
                    path_list = [layer_index]
                    node_tag = layer_index
                    node_identifier = str(current_depth) + "_" + str(layer_index)
                    tree.create_node(str(node_tag), node_identifier,
                                     parent="root",
                                     data={"value": grid.level_fisrt_node_value_noisy[layer_index], "path": path_list,
                                           "leaf": grid.level_fisrt_node_value_noisy[layer_index] < self.threshold})
                                     # data={"value": level_fisrt_data[layer_index], "path": path_list,
                                     #       "leaf": level_fisrt_data[layer_index] < self.threshold})
            else:
                layer_index = 0
                previous_depth = current_depth - 1
                previous_depth_nodes = self.get_previous_depth_nodes(tree, previous_depth)
                for parent_node in previous_depth_nodes:
                    subcell_index = int(parent_node.tag)
                    this_identifier = parent_node.identifier
                    this_neighbors = self.neighbors[subcell_index]
                    parent_path = parent_node.data["path"]
                    #结束节点
                    end_data = trajectory_set1.calculate_end_data(parent_path)
                    tree.create_node("&", str(current_depth) + "_" + str(layer_index),
                                     parent=parent_node.identifier,
                                     data={"value": end_data, "path": parent_path,
                                           "leaf": True})
                    layer_index += 1

                    for neighbor in self.neighbors[subcell_index]:
                        node_tag = neighbor
                        node_identifier = str(current_depth) + "_" + str(layer_index)
                        path_new_child = parent_path + [neighbor]

                        if len(path_new_child) != current_depth:
                            print("路径长度不等于当前深度")
                            print(node_tag,node_identifier)

                        path_array = np.array(path_new_child)
                        data_new_chlid = trajectory_set1.calculate_child_data(path_array, node_tag)
                        tree.create_node(str(node_tag), node_identifier,
                                         parent=parent_node.identifier,
                                         data={"value": data_new_chlid, "path": path_new_child,
                                               "leaf": data_new_chlid < self.threshold})
                        #如果是最后一层则必须为叶子节点
                        if current_depth == max_tree_depth:
                            tree.get_node(node_identifier).data["leaf"] = True

                        # child_node = tree.get_node(node_identifier)
                        # if child_node.data["value"] > parent_node.data["value"]:
                        #     print(child_node.tag + "的出现频率大于父节点" + parent_node.tag + "的出现频率")
                        layer_index += 1
                    #进行加噪
                    self.add_noise_for_children(tree, current_depth,parent_node.identifier,total_epsilon,grid,trajectory_set1)
                    self.consistency(tree,current_depth,parent_node.identifier)
                    # children = tree.children(parent_node.identifier)
            if self.only_leaf_nodes(tree, current_depth):
                last_depth_only_leaf_nodes = True
            if last_depth_only_leaf_nodes and current_depth < max_tree_depth:
                print("未达到最高层，最后一层节点均为叶子节点，停止构造树，最高层为：" + str(current_depth))
            current_depth += 1
        print("构造树完成")
        print(datetime.datetime.now())
        print("树的高度：",tree.depth())
        print("树的节点个数：",tree.size())
        # for node in tree.all_nodes():
        #     if tree.depth(node) == 2:
        #         print(node.data["leaf"])
        # print()
        # tree.show()
        self.tree = tree


        # for node in tree.expand_tree(mode=Tree.DEPTH):
        #     parent_node = tree.get_node(node)
        #     if parent_node is None:
        #         continue
        #     children = tree.children(node)  # 获取当前节点的所有子节点
        #     if not children:  # 如果没有子节点，跳过
        #         continue
        #     parent_value = parent_node.data['value']  # 父节点的值
        #     children_sum = sum(child.data['value'] for child in children)  # 子节点值的和
        #     if parent_value != children_sum:
        #         print(f"父节点 {parent_node.identifier} 的值 ({parent_value}) 不等于子节点值的和 ({children_sum})")

    def only_leaf_nodes(self, tree, last_depth):
        nodes_at_last_depth = []
        all_nodes = tree.all_nodes()
        for node in all_nodes:
            if tree.depth(node) == last_depth:
                nodes_at_last_depth.append(node)
        all_leaves = True
        for node in nodes_at_last_depth:
            if not node.data["leaf"]:
                all_leaves = False
                break
        return all_leaves

    def get_previous_depth_nodes(self, tree, previous_depth):
        previous_depth_nodes = []
        all_nodes = tree.all_nodes()
        for node in all_nodes:
            if tree.depth(node) == previous_depth and not node.data["leaf"]:
                previous_depth_nodes.append(node)
        return previous_depth_nodes

    def add_noise_for_children(self, tree, current_depth, parent_identifier, total_epsilon,grid : Grid,trajectory_set1: TrajectorySet):
        children = tree.children(parent_identifier)
        parent = tree.get_node(parent_identifier)
        child_values = [child.data["value"] for child in children]
        unpass_to_pass = False
        if sum(child_values) < self.threshold:
            self.false_nodes.append(parent_identifier)
            unpass_to_pass = True
        #计算衰减率
        siblings = tree.siblings(parent_identifier)
        siblings_values = [sibling.data["value"] for sibling in siblings]
        siblings_values.append(parent.data["value"])
        max_value = max(siblings_values)
        total_value = sum(siblings_values)
        decay_ratio = max_value / total_value
        parent_value =  parent.data["value"]
        tree_height = 0
        current_value = parent_value
        iterations = 0
        remaining_depth = self.max_depth - current_depth + 1
        sensitivity = 1
        noise1 = Noise()
        # 计算需要多少次衰减才会低于阈值
        while current_value >= self.threshold and iterations < self.max_depth:
            current_value *= decay_ratio
            tree_height += 1
            iterations += 1

        tree_height = min(tree_height, remaining_depth)
        if tree_height == 0:
            print("decay_ratio",decay_ratio)
            print("current_value",current_value)
            print("tree_height: ",tree_height)
            print("remaining_depth: ",remaining_depth)
        epsilon = total_epsilon / tree_height
        original_values = np.array([child.data["value"] for child in children])
        noisy_values = noise1.add_laplace_for_tree(original_values, epsilon, sensitivity)
        for child_node, noisy_value in zip(children, noisy_values):
            child_node.data["value"] = noisy_value
            child_node.data["leaf"] = noisy_value < self.threshold
            if child_node.tag == "&" or current_depth == self.max_depth:
                child_node.data["leaf"] = True

        noise_children = tree.children(parent_identifier)
        noise_children_value = [float(c.data["value"]) for c in noise_children]
        if math.isclose(sum(noise_children_value), 0, abs_tol=1e-12):
            self.children_all_zero_nodes.append(parent_identifier)
            print("加噪后子节点和为0：",parent_identifier)
            print("未加噪时节点值：",child_values)
            print("加噪后子节点值：",noise_children_value,noisy_values)
        # if unpass_to_pass:
        #     if current_depth == 2:
        #         if max(noise_children_value) == 0:
        #             print("父节点：",parent_identifier)
        #         print("父节点", parent.id0entifier, "没有子节点，父节点加噪值为：", parent.data["value"], "未加噪值：",
        #               grid.level_fisrt_node_value[int(parent.tag)], "是否叶子：", parent.data["leaf"])
        #         print("最大孩子：", max(noise_children_value), "孩子节点：", noise_children_value)
        #     else:
        #         if max(noise_children_value) == 0:
        #             print("父节点：",parent_identifier)
        #         print("父节点", parent.identifier, "没有子节点，父节点加噪值为：", parent.data["value"], "未加噪值：",
        #               trajectory_set1.calculate_child_data(
        #                   np.array(tree.parent(parent_identifier).data["path"] + [int(parent.tag)]),
        #                   int(parent.tag))
        #               , "是否叶子：", parent.data["leaf"])
        #         print("最大孩子：", max(noise_children_value), "孩子节点：", noise_children_value)


    def consistency(self, tree, current_depth, parent_identifier):
        parent = tree.get_node(parent_identifier)
        parent_value = parent.data["value"]
        children = tree.children(parent_identifier)
        child_values = [child.data["value"] for child in children]

        if all(v >= self.threshold for v in child_values):
            total = sum(child_values)
            for child in children:
                child.data["value"] = child.data["value"] / total * parent_value

        elif all(v < self.threshold for v in child_values):
            total = sum(child_values)
            if total > 0:
                for child in children:
                    child.data["value"] = child.data["value"] / total * parent_value
            else:
                for child in children:
                    child.data["value"] = 0
            # for child in children:
            #     child.data["value"] = 0

        else:
            # total = sum(child_values)
            # for child in children:
            #     child.data["value"] = child.data["value"] / total * parent_value
            above_threshold_nodes = [child for child in children if child.data["value"] >= self.threshold]
            below_threshold_nodes = [child for child in children if child.data["value"] < self.threshold]

            sum_above = sum(child.data["value"] for child in above_threshold_nodes)

            if sum_above > parent_value:
                for child in below_threshold_nodes:
                    child.data["value"] = 0
                for child in above_threshold_nodes:
                    child.data["value"] = child.data["value"] / sum_above * parent_value
            else:
                diff = parent_value - sum_above
                count_below = len(below_threshold_nodes)
                if count_below > 0:
                    add_value = diff / count_below
                    for child in below_threshold_nodes:
                        child.data["value"] = add_value
                total = sum(child.data["value"] for child in children)
                for child in children:
                    child.data["value"] = child.data["value"] / total * parent_value

        for child in children:
            child.data["leaf"] = child.data["value"] < self.threshold
            if child.tag == "&" or current_depth == self.max_depth:
                child.data["leaf"] = True

            if child.data["value"] > parent_value:
                print("子节点的值大于父节点的值，子节点：", child.tag, "子节点的值：", child.data["value"], "父节点的值：", parent_value)



